import pandas as pd
from openpyxl import load_workbook
from openpyxl.utils import get_column_letter

def fill_data_to_template(source_file, template_file, output_file, sheet_mapping=None, start_at_row=1):
    """
    将源表格中所有工作表的数据按列名填入模板表格，保持模板样式不变
    
    参数:
    source_file: 源Excel文件路径
    template_file: 模板Excel文件路径（包含样式）
    output_file: 输出Excel文件路径
    sheet_mapping: 工作表映射字典，格式为 {源工作表名: 模板工作表名}。
                   如果为None，则默认使用相同名称的工作表。
    start_at_row: 从源表格的第几行数据开始填入（1表示第一行数据）
    """
    def process_single_sheet(ws, source_df):
        """处理单个工作表的内部函数"""
        # 获取模板中的列名（假设第一行是标题）
        template_columns = []
        for cell in ws[1]:
            template_columns.append(cell.value)
        
        # 保存样式信息
        # 保存列宽
        column_widths = {}
        for col_letter in range(1, ws.max_column + 1):
            col_letter_str = get_column_letter(col_letter)
            column_widths[col_letter_str] = ws.column_dimensions[col_letter_str].width
        
        # 保存行高（标题行）
        header_row_height = ws.row_dimensions[1].height
        
        # 清除模板中除标题外的数据，但保留样式
        max_row = ws.max_row
        for row_idx in range(2, max_row + 1):
            for col_idx in range(1, ws.max_column + 1):
                cell = ws.cell(row=row_idx, column=col_idx)
                cell.value = None
        
        # 按列名匹配并填充数据
        for idx, row in source_df.iterrows():
            for col_idx, col_name in enumerate(template_columns, 1):
                if col_name in source_df.columns:
                    target_row = idx + 2
                    cell = ws.cell(row=target_row, column=col_idx)
                    cell.value = row[col_name]
        
        # 恢复列宽
        for col_letter, width in column_widths.items():
            ws.column_dimensions[col_letter].width = width
        
        # 恢复标题行高
        ws.row_dimensions[1].height = header_row_height
        
        # 确定最后一行（修复空数据时的潜在bug）
        last_row = 1  # 至少有一行标题
        if not source_df.empty:
            last_row = len(source_df) + 1
        
        # 调整数据行高
        for row_idx in range(2, last_row + 1):
            ws.row_dimensions[row_idx].height = header_row_height
    
    # 读取源数据的所有工作表
    try:
        source_sheets = pd.read_excel(source_file, sheet_name=None, engine='openpyxl')
    except Exception as e:
        raise ValueError(f"无法读取源文件 {source_file}: {str(e)}\n"
                         "可能原因：1. 文件格式不兼容 2. 文件损坏 3. openpyxl版本问题\n"
                         "建议：1. 尝试将文件另存为新的.xlsx文件 2. 更新openpyxl库 (pip install --upgrade openpyxl)")
    
    # 加载模板工作簿（保留样式）
    try:
        wb = load_workbook(template_file)
    except Exception as e:
        raise ValueError(f"无法加载模板文件 {template_file}: {str(e)}")
    
    # 如果提供了sheet_mapping，则使用它；否则，假设源和模板的工作表名称相同
    if sheet_mapping is None:
        # 只处理源和模板中都存在的工作表
        common_sheets = set(source_sheets.keys()) & set(wb.sheetnames)
        sheet_mapping = {sheet: sheet for sheet in common_sheets}
    
    # 处理每个工作表
    for src_sheet, tgt_sheet in sheet_mapping.items():
        if tgt_sheet not in wb.sheetnames:
            print(f"警告: 模板中不存在工作表 {tgt_sheet}，跳过源工作表 {src_sheet}")
            continue
        
        ws = wb[tgt_sheet]
        source_df = source_sheets[src_sheet]
        
        # 应用start_at_row过滤并重置索引
        start_at_row = max(1, start_at_row)  # 确保start_at_row至少为1
        if start_at_row > len(source_df):
            # 源表数据行数小于参数，不填入数据，设置为空DataFrame
            source_df = pd.DataFrame(columns=source_df.columns)
        else:
            pandas_start_idx = start_at_row - 1
            if pandas_start_idx > 0:
                source_df = source_df.iloc[pandas_start_idx:].reset_index(drop=True)
        
        process_single_sheet(ws, source_df)
    
    # 保存结果
    try:
        wb.save(output_file)
    except Exception as e:
        raise IOError(f"无法保存到 {output_file}: {str(e)}")

# 示例用法
if __name__ == "__main__":
    source_file = r"D:\wxzx\45g新建站模板11月.xlsx"
    template_file = r"D:\wxzx\数据修改模板-1107"
    output_file = r"D:\wxzx\output.xlsx"
    
    # 处理所有匹配的工作表，从第3行数据开始
    try:
        fill_data_to_template(source_file, template_file, output_file, start_at_row=3)
        print(f"数据已成功填入模板并保存到 {output_file}")
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
    
    # # 也可以指定特定工作表映射
    # sheet_mapping = {
    #     "AAU数据": "AAU",
    #     "RRU数据": "RRU"
    # }
    # try:
    #     fill_data_to_template(source_file, template_file, output_file, sheet_mapping)
    #     print(f"数据已成功填入模板并保存到 {output_file}")
    # except Exception as e:
    #     print(f"处理过程中出错: {str(e)}")